#include "stdafx.h"
#include "vcclr.h"
#include "AviFileReader.h"

using namespace System;

namespace ManagedVC1
{
    using namespace Exceptions;

    AviFileReader::AviFileReader() :
        m_aviFileInInfo(new AVIFILEINFO),
        m_videoInfoHeader(new BITMAPINFOHEADER),
        m_audioInfoHeader(new WAVEFORMATEX)
    {
        m_aviFileIn = nullptr;
        m_videoStream = nullptr;
        m_audioStream = nullptr;

        m_SwapUV = false;
        m_numFrames = 0;
        m_YSample = 0;
        m_UVSample = 0;
        m_audioSamplesPos = 0;
        m_firstVideoFrame = 0;
        m_firstAudioFrame = 0;
        m_lastAudioFrame = 0;
        m_currentFrameNumber = 0;
        m_audioSamplesPerFrame = 0;
        m_audioBufferSize = 0;

        ZeroMemory(m_aviFileInInfo, sizeof(AVIFILEINFO));
        ZeroMemory(m_videoInfoHeader, sizeof(BITMAPINFOHEADER));
        ZeroMemory(m_audioInfoHeader, sizeof(WAVEFORMATEX));
    }

    void AviFileReader::Initialize(String^ fileName, CommonFrameRate^ frameRate)
    {
        if (String::IsNullOrEmpty(fileName) || !System::IO::File::Exists(fileName)) throw gcnew ArgumentException("You must specify a valid file", "fileName");
        m_frameRate = frameRate;

        long hr = 0;

        AVIFileInit();

        // open the AVI file
        // pointer is on the stack, so it won't move around
        pin_ptr<PAVIFILE> pavi = &m_aviFileIn;
        pin_ptr<const wchar_t> wfileName = PtrToStringChars(fileName);
        hr = AVIFileOpen(pavi, wfileName, OF_READ, nullptr);
        ExceptionUtility::ThrowExceptionForHr(hr);

        // get information on the AVI file
        hr = AVIFileInfo(m_aviFileIn, m_aviFileInInfo, sizeof(AVIFILEINFO));
        ExceptionUtility::ThrowExceptionForHr(hr);

        // iterate over all the streams
        int numStreams = m_aviFileInInfo->dwStreams;
        for (int i = 0; i < numStreams; i++)
        {
            IAVIStream* aviStream = nullptr;
            AutoPtr<AVISTREAMINFO> streamInfo(new AVISTREAMINFO);

            // get the stream
            hr = AVIFileGetStream(m_aviFileIn, &aviStream, 0, i);
            ExceptionUtility::ThrowExceptionForHr(hr);

            // Get the information for the stream
            hr = AVIStreamInfo(aviStream, streamInfo, sizeof(AVISTREAMINFO));
            ExceptionUtility::ThrowExceptionForHr(hr);

            if (streamInfo->fccType == streamtypeVIDEO)
            {
                if (streamInfo->fccHandler == FOURCC_YV12)
                {
                    m_SwapUV = true;
                    streamInfo->fccHandler = FOURCC_I420;
                }
                else if (streamInfo->fccHandler != FOURCC_I420 &&
                    streamInfo->fccHandler != FOURCC_IYUV &&
                    streamInfo->fccHandler != FOURCC_iyuv &&
                    streamInfo->fccHandler != FOURCC_i420)
                {
                    Exception^ exp = gcnew Exception("Bad AVS/AVI format - Has to be YV12!");
                    exp->Data["fccHandler"] = streamInfo->fccHandler;
                    throw exp;
                }

                // reassign stream pointers
                m_videoStream = aviStream;
                m_videoStreamInfo.Reset(streamInfo.Release());

                // get basic stream information
                m_numFrames = m_videoStreamInfo->dwLength;
                m_firstVideoFrame = AVIStreamStart(m_videoStream);

                // get size of format data
                long videoFormatSize = 0;
                pin_ptr<long> videoSizePtr = &videoFormatSize;
                hr = AVIStreamReadFormat(m_videoStream, 0, nullptr, videoSizePtr);
                ExceptionUtility::ThrowExceptionForHr(hr);

                // query the format data
                hr = AVIStreamReadFormat(m_videoStream, 0, m_videoInfoHeader, videoSizePtr);
                ExceptionUtility::ThrowExceptionForHr(hr);

                m_YSample = m_videoInfoHeader->biWidth * m_videoInfoHeader->biHeight;
                m_UVSample = m_YSample / 4;
            }
            else if (streamInfo->fccType == streamtypeAUDIO)
            {
                // reassign pointers
                m_audioStream = aviStream;
                m_audioStreamInfo.Reset(streamInfo.Release());

                // get basic stream info
                m_firstAudioFrame = AVIStreamStart(m_audioStream);
                m_lastAudioFrame = m_firstAudioFrame + AVIStreamLength(m_audioStream);

                long audioFormatSize = 0;
                pin_ptr<long> audioSizePtr = &audioFormatSize;
                hr = AVIStreamReadFormat(m_audioStream, 0, nullptr, audioSizePtr);
                ExceptionUtility::ThrowExceptionForHr(hr);
                
                // read audio format data
                hr = AVIStreamReadFormat(m_audioStream, 0, m_audioInfoHeader, audioSizePtr);
                ExceptionUtility::ThrowExceptionForHr(hr);

                // calculate size information
                m_audioSamplesPerFrame = (int) (m_audioInfoHeader->nSamplesPerSec / m_frameRate->Rate);
                m_audioBufferSize = (int) ((m_audioSamplesPerFrame) * (m_audioInfoHeader->wBitsPerSample * m_audioInfoHeader->nChannels / 8));

                // allocate and zero the audio buffer
                m_audioBuffer.Reset(new BYTE[m_audioBufferSize]);
                if (m_audioBuffer.GetPointer() == nullptr) throw gcnew OutOfMemoryException();
                ZeroMemory(m_audioBuffer, m_audioBufferSize);
            }
			else
			{
				AVIStreamRelease(aviStream);
			}
        }
    }

    __int64 AviFileReader::ReadVideoFrame(AutoArrayPtr<BYTE>% pVideo)
    {
        if (m_videoStream == nullptr) { pVideo.Reset(); return 0; }
        
        long hr = 0;
        long cbSample = 0, cSamples = 0, bytesRead = 0;

        // get the necessary size to store the video frame
        hr = AVIStreamRead(m_videoStream, m_currentFrameNumber, AVISTREAMREAD_CONVENIENT, nullptr, 0, &cbSample, &cSamples);
        ExceptionUtility::ThrowExceptionForHr(hr);

        // clear existing memory and allocate new
        m_videoBuffer.Reset(new BYTE[cbSample]);
        if (m_videoBuffer.GetPointer() == nullptr) throw gcnew OutOfMemoryException();
        ZeroMemory(m_videoBuffer, cbSample);

        // fill the video buffer
        hr = AVIStreamRead(m_videoStream, m_currentFrameNumber, AVISTREAMREAD_CONVENIENT, m_videoBuffer, cbSample, &bytesRead, nullptr);
        ExceptionUtility::ThrowExceptionForHr(hr);

        if (m_SwapUV)
        {
            AutoArrayPtr<BYTE> tempVideoBuffer(new BYTE[cbSample]);
            ZeroMemory(tempVideoBuffer, cbSample);

            // swap the U and V sections in memory
            memcpy(tempVideoBuffer, m_videoBuffer, m_YSample);
            memcpy(tempVideoBuffer + m_YSample + m_UVSample, m_videoBuffer + m_YSample, m_UVSample);
            memcpy(tempVideoBuffer + m_YSample, m_videoBuffer + m_YSample + m_UVSample, m_UVSample);
            m_videoBuffer.Reset(tempVideoBuffer.Release());
        }

        pVideo.Reset(m_videoBuffer.Release());
        return GetVideoTimestamp(m_currentFrameNumber++);
    }

    __int64 AviFileReader::GetVideoTimestamp(int frameNumber)
    {
        __int64 ts = 0;
        if (frameNumber > 0)
        {
            ts = (__int64)((m_frameRate->Denominator * (double)frameNumber * 10000000.0) / (m_frameRate->Numerator) + 0.5);
        }
        return ts;
    }

    __int64 AviFileReader::ReadAudioFrame(BYTE** pAudio, long* audioLen)
    {
        if (audioLen == nullptr || pAudio == nullptr) throw gcnew NullReferenceException();
        if (m_audioStream == nullptr) { pAudio = nullptr; *audioLen = 0; return 0; };

        long hr = 0;
        long bytesRead = 0;

        hr = AVIStreamRead(m_audioStream, m_audioSamplesPos, m_audioSamplesPerFrame, m_audioBuffer, m_audioBufferSize, &bytesRead, nullptr);

        if (SUCCEEDED(hr) && m_audioSamplesPos <= m_lastAudioFrame)
        {
            *pAudio = m_audioBuffer;
            *audioLen = bytesRead;

            m_audioSamplesPos += m_audioSamplesPerFrame;
            return GetAudioTimestamp();
        }
        else if (hr == AVIERR_FILEREAD || hr == AVIERR_NODATA || m_audioSamplesPos > m_lastAudioFrame)
        {
            *audioLen = 0;
            return 0;
        }
        else ExceptionUtility::ThrowExceptionForHr(hr, "AVIStreamRead");

        return 0;
    }

    __int64 AviFileReader::GetAudioTimestamp()
    {
        __int64 ts = 0;

        if (m_audioSamplesPos != 0)
        {
            ts = (__int64)((m_audioSamplesPos - m_audioSamplesPerFrame) / (m_audioInfoHeader->nSamplesPerSec * 1.0) * 10000000.0);
        }

        return ts;
    }

    void AviFileReader::SeekToBeginning()
    {
        m_currentFrameNumber = m_firstVideoFrame;
        m_audioSamplesPos = m_firstAudioFrame;
    }
}